import inspect
import os
import stat
import numpy as np
import torch
import re
import csv

try:
    import torch_npu
except ImportError:
    is_gpu=True
else:
    is_gpu=False

from ..common.utils import check_file_or_directory_path, print_error_log, __version__, \
    print_warn_log, CompareException, Const, get_time, print_info_log, modify_dump_path


DumpCount = 0
forward_init_status = False
backward_init_status = False
range_begin_flag, range_end_flag = False, False


class DumpUtil(object):
    dump_data_dir = None
    dump_path = None
    dump_switch = None
    dump_switch_mode = Const.ALL
    dump_switch_scope = []
    dump_init_enable = False
    dump_api_list = []
    backward_input = {}
    dump_dir_tag = 'ptdbg_dump'

    @staticmethod
    def set_dump_path(save_path):
        DumpUtil.dump_path = save_path
        DumpUtil.dump_init_enable = True

    @staticmethod
    def set_dump_switch(switch, mode, scope, api_list):
        DumpUtil.dump_switch = switch
        DumpUtil.dump_switch_mode = mode
        DumpUtil.dump_init_enable = True
        DumpUtil.dump_switch_scope = scope
        DumpUtil.dump_api_list = [api.lower() for api in api_list]
        if mode == Const.ACL:
            DumpUtil.dump_switch_scope = [api_name.replace("backward", "forward") for api_name in scope]

    def check_list_or_acl_mode(name_prefix):
        global DumpCount
        for item in DumpUtil.dump_switch_scope:
            if name_prefix.startswith(item):
                DumpCount = DumpCount + 1
                return True

    def check_range_mode(name_prefix):
        global range_begin_flag
        global range_end_flag
        if name_prefix.startswith(DumpUtil.dump_switch_scope[0]): 
            range_begin_flag = True
            return True
        if name_prefix.startswith(DumpUtil.dump_switch_scope[1]):
            range_end_flag = True
            return True
        if range_begin_flag and not range_end_flag:
            return True
        return False

    def check_stack_mode(name_prefix):
        if len(DumpUtil.dump_switch_scope) == 0:
            return True
        elif len(DumpUtil.dump_switch_scope) == 1:
            return name_prefix.startswith(DumpUtil.dump_switch_scope[0])
        else:
            print_info_log("dump scope is invalid, Please set the scope mode in"
                            " set_dump_switch with 'all', 'list', 'range', 'stack', 'acl', 'api_list'!") 
        return False

    check_mapper = {
        Const.LIST: check_list_or_acl_mode,
        Const.ACL: check_list_or_acl_mode,
        Const.RANGE: check_range_mode,
        Const.STACK: check_stack_mode
    }

    @staticmethod
    def check_switch_scope(name_prefix):
        if DumpUtil.dump_switch_mode in DumpUtil.check_mapper:
            check_func = DumpUtil.check_mapper[DumpUtil.dump_switch_mode]
            return check_func(name_prefix)
        return False

    @staticmethod
    def get_dump_path():
        if DumpUtil.dump_path:
            return DumpUtil.dump_path

        if DumpUtil.dump_switch_mode == Const.ALL:
            raise RuntimeError("get_dump_path: the file path is empty,"
                               " you must use set_dump_path to set a valid dump path!!!")
        else:
            dir_path = os.path.realpath("./")
            dump_file_name = "scope_dump_{}_{}_{}.pkl".format(
                DumpUtil.dump_switch_mode, DumpUtil.dump_switch_scope[0], get_time())
            DumpUtil.dump_path = os.path.join(dir_path, dump_file_name)
            return DumpUtil.dump_path

    @staticmethod
    def get_dump_switch():
        if DumpUtil.dump_switch is None:
            return False
        return DumpUtil.dump_switch == "ON"


def set_dump_path(fpath=None, dump_tag='ptdbg_dump'): 
    if fpath is None:
        raise RuntimeError("set_dump_path '{}' error, please set a valid filename".format(fpath))
        return
    real_path = os.path.realpath(fpath)
    if os.path.isdir(real_path):
        print_error_log("set_dump_path '{}' error, please set a valid filename.".format(real_path))
        raise CompareException(CompareException.INVALID_PATH_ERROR)
    if os.path.exists(real_path):
        os.remove(real_path)
    DumpUtil.set_dump_path(real_path)
    DumpUtil.dump_dir_tag = dump_tag


def set_dump_switch(switch, mode=Const.ALL, scope=[], api_list=[]): 
    global DumpCount
    if not re.match(r"^O(N|FF)(.*)*$", switch): 
        print_error_log("Please set dump switch with 'ON' or 'OFF'.")
    if mode == Const.LIST and switch == "ON":
        DumpCount = 0
    if mode == Const.LIST and switch == "OFF":
        print_info_log("The number of matched dump is {}".format(DumpCount))
    if mode == Const.RANGE:
        assert len(scope) == 2, "set_dump_switch, scope param set invalid, it's must be [start, end]."
    if mode == Const.LIST:
        assert len(scope) != 0, "set_dump_switch, scope param set invalid, it's should not be an empty list."
    if mode == Const.STACK:
        assert len(scope) <= 2, "set_dump_switch, scope param set invalid, it's must be [start, end] or []."
    if mode == Const.ACL:
        assert len(scope) == 1, "set_dump_switch, scope param set invalid, only one api name is supported in acl mode."
    DumpUtil.set_dump_switch(switch, mode=mode, scope=scope, api_list=api_list)


def _dump_tensor_completely(x, prefix, dump_file_name):
    dump_flag = Const.DUMP_RATIO_MAX + 1
    if isinstance(x, (tuple, list)) and x:
        for i, item in enumerate(x):
            _dump_tensor_completely(item, "{}.{}".format(prefix, i), dump_file_name) 
    else:
        with os.fdopen(os.open(dump_file_name, os.O_RDWR | os.O_CREAT, stat.S_IWUSR | stat.S_IRUSR), "a") as f:
            if isinstance(x, torch.Tensor) and x.numel() != 0:
                output_path = os.path.join(DumpUtil.dump_data_dir, f'{prefix}.npy')
                save_tensor = x.contiguous().cpu().detach().numpy()
                np.save(output_path, save_tensor)
                writer = csv.writer(f)
                writer.writerows([prefix, dump_flag, [], str(x.dtype), tuple(x.shape)]) 
    


def dump_stack_info(name_template, dump_file):
    stack_str = []
    for (_, path, line, func, code, _) in inspect.stack()[3:]:
        stack_line = [path, str(line), func, code[0].strip() if code else code]
        stack_str.append(stack_line)
    _dump_tensor_completely(stack_str, name_template.format("stack_info"), dump_file)


def cosine_similarity(n_value, b_value):
    np.seterr(divide='ignore', invalid='ignore')
    num = n_value.dot(b_value)
    a_norm = np.linalg.norm(n_value)
    b_norm = np.linalg.norm(b_value)
    message = ''
    cos = num / (a_norm * b_norm) 
    if np.isnan(cos):
        message = 'Cannot compare by Cosine Similarity, the dump data has NaN.'
        result = Const.NAN
    return result, message


def get_max_abs_err(n_value, b_value):
    temp_res = n_value - b_value
    max_value = np.max(np.abs(temp_res))
    return max_value, ""


def compare_by_op(op_name, op_name_mapping_dict, input_parma):
    npu_bench_name_list = op_name_mapping_dict[op_name]
    try:
        n_value = np.load(os.path.join(input_parma.get("npu_dump_data_dir"), npu_bench_name_list[0] + ".npy"), allow_pickle=True) 
        b_value = np.load(os.path.join(input_parma.get("bench_dump_data_dir"), npu_bench_name_list[1] + ".npy"), allow_pickle=True)
        n_value = n_value.reshape(-1).astype(float)
        b_value = b_value.reshape(-1).astype(float)
    except IOError as error:
        return " ", "", "Dump file:{} not found".format(error.filename)
    err_msg = ""
    cos_sim, message = cosine_similarity(n_value, b_value)
    err_msg += message
    max_abs_err, _ = get_max_abs_err(n_value, b_value)
    return cos_sim, max_abs_err, err_msg